from vquant.store import Store


class Position(Store):
    columns = ['datetime', 'symbol', 'cost', 'quantity', 'margin', 'profit']

    def get_or_create(self, datetime, symbol):
        positions = self.query(symbol=dict(eq=symbol))
        if positions.shape[0]:
            return positions.iloc[0]
        return self.insert(datetime=datetime, symbol=symbol, quantity=0, cost=0, margin=0, profit=0)

    def replace(self, position, price, quantity_multiple):
        position.profit = price * position.quantity * quantity_multiple - position.cost
        self.update(position.to_dict(), symbol=dict(eq=position.symbol))
        return position

    def add(self, datetime, symbol, price, quantity, leverage, quantity_multiple):
        opened, closed, profit = quantity, 0, 0
        position = self.get_or_create(datetime, symbol).copy()
        position = self.replace(position, price, quantity_multiple)
        holdings = position.quantity + quantity
        if position.quantity > 0:
            if quantity < 0:
                closed = -min(abs(quantity), position.quantity)
                opened = holdings if holdings < 0 else 0
        elif position.quantity < 0:
            if quantity > 0:
                closed = min(abs(quantity), position.quantity)
                opened = holdings if holdings > 0 else 0
        if closed:
            avg_profit = position.profit / abs(position.quantity)
            position.profit = avg_profit * holdings
            position.margin = position.margin / abs(position.quantity) * holdings
            position.cost = position.cost / position.quantity * holdings
            profit = avg_profit * abs(closed)
        if opened:
            position.cost = position.cost + price * opened * quantity_multiple
            position.margin = position.margin + abs(price * opened) * quantity_multiple / leverage
        position.datetime = datetime
        position.quantity = holdings
        self.update(position.to_dict(), symbol=dict(eq=position.symbol))
        return position, profit


__all__ = [Position]
